# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import functools
import json
import logging

from tornado import gen
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect, WebSocketClosedError
from tornado.concurrent import Future

from coin.base.timestamp import get_timestamp
from coin.exchange.base.rest.task_scheduler import Scheduler

BITSTAMP_WS_API_URL = 'wss://ws.bitstamp.net'


class ChannelType(enum.Enum):
  TRADE = 'live_trades'
  ORDERS = 'live_orders'
  ORDER_BOOK = 'order_book'
  DETAILED_ORDER_BOOK = 'detailed_order_book'
  FULL_ORDER_BOOK = 'diff_order_book'


class EventType(enum.Enum):
  SUBSCRIBE = 'bts:subscribe'
  UNSUBSCRIBE = 'bts:unsubscribe'
  ERROR = 'bts:error'
  SUBSCRIPTION_SUCCEEDED = 'bts:subscription_succeeded'
  REQUEST_RECONNECT = 'bts:request_reconnect'
  DATA = 'data'
  TRADE = 'trade'
  ORDER_CREATED = 'order_created'
  ORDER_CHANGED = 'order_changed'
  ORDER_DELETED = 'order_deleted'


def get_channel_name(channel_type: ChannelType, currency_pair):
  return '%s_%s' % (channel_type.value, currency_pair)


@functools.lru_cache(maxsize=64)
def parse_channel_name(channel_name):
  for channel in ChannelType:
    if channel_name.startswith(channel.value):
      return {'channel_type': channel, 'currency_pair': channel_name[len(channel.value) + 1:]}
  raise ValueError('invalid channel name %s' % channel_name)


class BitstampWsPrivateClient(object):
  def __init__(
      self,
      *,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      on_ready=None,
      timeout_sec=10,
      api_url=BITSTAMP_WS_API_URL,
  ):

    self._ioloop = ioloop or IOLoop.current()
    self._api_url = api_url
    self._ws = None
    self._ready = False
    self._timeout = timeout_sec * 1e9
    self._num_timeouts = 0

    self._notification_callback = notification_callback or (lambda x: None)
    self._close_callback = close_callback or (lambda: None)
    self._on_ready = on_ready or (lambda: None)
    self._scheduler = Scheduler()

    self._last_received_message_timestamp = None
    self._pending_future = None
    self._subscribed_channels = set()

    self._stopped = False
    self._init_params = (ioloop, notification_callback, close_callback, timeout_sec)
    self._logger = logging.getLogger(__name__)

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _check_future_timeout(self):
    timestamp = get_timestamp()
    fut = self._pending_future
    if fut is None:
      return
    if timestamp - fut.timestamp > self._timeout:
      self._num_timeouts += 1
      self._pending_future = None
      fut.set_exception(TimeoutError('Websocket request timeout! %s)' % fut.method))

  def _check_connection(self):
    if self._last_received_message_timestamp is None:
      return

    timestamp = get_timestamp()
    diff_time = timestamp - self._last_received_message_timestamp
    if diff_time > self._timeout:
      self._logger.warning('No message received for %s seconds', diff_time / 1e9)
      self.stop()

  def _write_message(self, msg):
    try:
      self._ws.write_message(json.dumps(msg))
    except WebSocketClosedError as e:
      self._logger.error('Write message error, ws is close!')
      self._pending_future.set_exception(e)
      self._pending_future = None

  def log_json_message(self, msg):
    self._logger.info('message: %s\n', json.dumps(msg, indent=2))

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket.', type(self).__name__)
      self._ws = yield websocket_connect(self._api_url)
      self._ioloop.add_callback(self._run)
      self._ready = True
      self._on_ready()
      self._logger.info('Bitstamp private client websocket connected.')
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def on_message(self, raw_msg: bytes):
    timestamp = get_timestamp()
    self._last_received_message_timestamp = timestamp
    if self._stopped:
      return

    try:
      data = json.loads(raw_msg)
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return

    event = EventType(data['event'])
    if event == EventType.ERROR:
      self._logger.error('Subscription error!')
      if self._pending_future is not None:
        self._pending_future.set_exception(ValueError(json.dumps(data)))
        self._pending_future = None
      else:
        self._logger.error('Unexpected msg received! %s', data)
      self.stop()
    elif event == EventType.REQUEST_RECONNECT:
      self._logger.error('Exchange server request reconnect! %s', data)
      self.stop()
    elif event == EventType.SUBSCRIPTION_SUCCEEDED:
      if self._pending_future is not None:
        self._subscribed_channels.add(data['channel'])
        self._pending_future.set_result(data)
        self._pending_future = None
      else:
        self._logger.error('Unexpected msg received! %s', data)
    else:
      self._notification_callback(data)

  def send_request(self, msg):
    if self._ws is None:
      self._logger.error('websocket not connected yet! %s', msg)
      return

    if self._pending_future is not None:
      self._logger.error('Previous futures is not resolved yet!')
      return

    fut = Future()
    fut.timestamp = get_timestamp()
    self._pending_future = fut
    self.log_json_message(msg)
    self._ioloop.add_callback(self._write_message, msg)
    return fut

  def subscribe_channel(self, channel_type: ChannelType, currency_pair):
    channel_name = get_channel_name(channel_type, currency_pair)
    msg = {'event': EventType.SUBSCRIBE.value, 'data': {'channel': channel_name}}
    fut = self.send_request(msg)
    return fut

  def unsubscribe_channel(self, channel_type: ChannelType, currency_pair):
    channel_name = get_channel_name(channel_type, currency_pair)
    msg = {'event': EventType.UNSUBSCRIBE.value, 'data': {'channel': channel_name}}
    fut = self.send_request(msg)
    return fut

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._scheduler.add_periodic_task(1, self._check_future_timeout)
    self._scheduler.add_periodic_task(10, self._check_connection)
    self._scheduler.start()

  def stop(self):
    if self._stopped:
      return
    self._scheduler.stop()
    self._ready = False
    self._ioloop.add_callback(self._async_stop)
    self._logger.info('Bitstamp ws private client (%s) stopped!', id(self))

  @gen.coroutine
  def _async_stop(self):
    for channel in self._subscribed_channels:
      params = parse_channel_name(channel)
      response = yield self.unsubscribe_channel(**params)
      self._logger.info('Topic (%s) unsubscribed!\n%s', channel, response)

    self._stopped = True
    if self._ws is not None:
      self._ws.close()
      self._ws = None
    self._close_callback()

  def clone(self):
    return type(self)(*self._init_params)
